import cv2
import numpy as np

def main():
    # 读取图像
    img1 = cv2.imread('10/5.png', cv2.IMREAD_GRAYSCALE)  # 查询图像
    img2 = cv2.imread('10/4.png', cv2.IMREAD_GRAYSCALE)  # 训练图像
    
    if img1 is not None and img2 is not None:
        # 初始化 SIFT 检测器
        sift = cv2.SIFT_create()
        
        # 检测关键点和计算描述符
        kp1, des1 = sift.detectAndCompute(img1, None)
        kp2, des2 = sift.detectAndCompute(img2, None)
        
        # 使用 FLANN 匹配器进行特征匹配
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        flann = cv2.FlannBasedMatcher(index_params, search_params)
        matches = flann.knnMatch(des1, des2, k=2)
        
        # 筛选优质匹配（Lowe's ratio test）
        good_matches = []
        for m, n in matches:
            if m.distance < 0.7 * n.distance:
                good_matches.append(m)
        
        # 提取匹配点的坐标
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        
        # 计算透视变换矩阵（使用 RANSAC 过滤异常值）
        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        
        # 应用透视变换
        height, width = img2.shape
        warped_img1 = cv2.warpPerspective(img1, M, (width, height))
        
        # 计算差分图
        diff = cv2.absdiff(warped_img1, img2)
        
        # 对差分图进行阈值处理
        _, thresholded = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY)
        
        # 使用中值滤波去除噪声
        thresholded = cv2.medianBlur(thresholded, 5)
        
        # 查找轮廓
        contours, _ = cv2.findContours(thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 在原图上绘制红色矩形框标记差异区域
        img2_color = cv2.cvtColor(img2, cv2.COLOR_GRAY2BGR)
        diff_color = cv2.cvtColor(diff, cv2.COLOR_GRAY2BGR)
        
        for contour in contours:
            # 忽略太小的区域
            if cv2.contourArea(contour) < 100:
                continue
            # 获取边界矩形
            x, y, w, h = cv2.boundingRect(contour)
            # 在原图和差分图上绘制红色矩形
            cv2.rectangle(img2_color, (x, y), (x+w, y+h), (0, 0, 255), 2)
            cv2.rectangle(diff_color, (x, y), (x+w, y+h), (0, 0, 255), 2)
        
        # 可视化关键点匹配
        matched_img = cv2.drawMatches(img1, kp1, img2, kp2, good_matches, None,
                                     flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
        
        # 显示结果
        cv2.imshow('Original Image 1', img1)
        cv2.imshow('Original Image 2', img2_color)
        cv2.imshow('Feature Matches', matched_img)
        cv2.imshow('Warped Image 1', warped_img1)
        cv2.imshow('Difference Map', diff_color)
        cv2.imshow('Thresholded Difference', thresholded)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("Error: Failed to load images!")

if __name__ == '__main__':
    main()